{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset, Dataset\n",
    "\n",
    "# 提示词模板\n",
    "SYSTEM_PROMPT = \"\"\"\n",
    "Respond in the following format:\n",
    "<reasoning>\n",
    "...\n",
    "</reasoning>\n",
    "<answer>\n",
    "...\n",
    "</answer>\n",
    "\"\"\"\n",
    "# 根据数据集文本格式提取答案\n",
    "def extract_hash_answer(text: str) -> str | None:\n",
    "    if \"####\" not in text:\n",
    "        return None\n",
    "    return text.split(\"####\")[1].strip()\n",
    "\n",
    "dataset = load_dataset(\"./gsm8k/\")['train']\n",
    "dataset = dataset.map(\n",
    "    lambda x: {\n",
    "        \"prompt\": [\n",
    "            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "            {\"role\": \"user\", \"content\": x[\"question\"]},\n",
    "        ],\n",
    "        \"answer\": extract_hash_answer(x[\"answer\"]),\n",
    "    }\n",
    ")\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 模型加载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsloth import FastLanguageModel,PatchFastRL ,is_bfloat16_supported\n",
    "\n",
    "max_seq_length = 1024 # 增加模型推理的长度\n",
    "lora_rank = 64 # LoRA的秩 可以是 8, 16, 32, 64, 128 该秩越大模型表现得越聪明 但速度会更慢\n",
    "\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name = \"./Qwen2.5-7B-Instruct\", #TODO :修改为你自己的模型位置\n",
    "    max_seq_length = max_seq_length,\n",
    "    load_in_4bit = True,\n",
    "    fast_inference = True, # 增加模型的推理速度\n",
    "    max_lora_rank = lora_rank,\n",
    "    gpu_memory_utilization = 0.5,\n",
    ")\n",
    "\n",
    "model = FastLanguageModel.get_peft_model(\n",
    "    model,\n",
    "    r = lora_rank,\n",
    "    target_modules = [\n",
    "        \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
    "        \"gate_proj\", \"up_proj\", \"down_proj\",\n",
    "    ], # 选择要进行LoRA的线性层\n",
    "    lora_alpha = lora_rank,\n",
    "    use_gradient_checkpointing = \"unsloth\", # 启用长上下文微调\n",
    "    random_state = 3407 # 控制随机数 保证结果可复现\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 设置奖励函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "XML_COT_FORMAT = \"\"\"\\\n",
    "<reasoning>\n",
    "{reasoning}\n",
    "</reasoning>\n",
    "<answer>\n",
    "{answer}\n",
    "</answer>\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "def extract_xml_answer(text: str) -> str:\n",
    "    answer = text.split(\"<answer>\")[-1]\n",
    "    answer = answer.split(\"</answer>\")[0]\n",
    "    return answer.strip()\n",
    "\n",
    "\n",
    "def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:\n",
    "    responses = [completion[0][\"content\"] for completion in completions]\n",
    "    q = prompts[0][-1][\"content\"]\n",
    "    extracted_responses = [extract_xml_answer(r) for r in responses]\n",
    "    print(\n",
    "        \"-\" * 20,\n",
    "        f\"Question:\\n{q}\",\n",
    "        f\"\\nAnswer:\\n{answer[0]}\",\n",
    "        f\"\\nResponse:\\n{responses[0]}\",\n",
    "        f\"\\nExtracted:\\n{extracted_responses[0]}\",\n",
    "    )\n",
    "    return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]\n",
    "\n",
    "\n",
    "def int_reward_func(completions, **kwargs) -> list[float]:\n",
    "    responses = [completion[0][\"content\"] for completion in completions]\n",
    "    extracted_responses = [extract_xml_answer(r) for r in responses]\n",
    "    return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]\n",
    "\n",
    "\n",
    "def strict_format_reward_func(completions, **kwargs) -> list[float]:\n",
    "    \"\"\"Reward function that checks if the completion has a specific format.\"\"\"\n",
    "    pattern = r\"^<reasoning>\\n.*?\\n</reasoning>\\n<answer>\\n.*?\\n</answer>\\n$\"\n",
    "    responses = [completion[0][\"content\"] for completion in completions]\n",
    "    matches = [re.match(pattern, r) for r in responses]\n",
    "    return [0.5 if match else 0.0 for match in matches]\n",
    "\n",
    "\n",
    "def soft_format_reward_func(completions, **kwargs) -> list[float]:\n",
    "    \"\"\"Reward function that checks if the completion has a specific format.\"\"\"\n",
    "    pattern = r\"<reasoning>.*?</reasoning>\\s*<answer>.*?</answer>\"\n",
    "    responses = [completion[0][\"content\"] for completion in completions]\n",
    "    matches = [re.match(pattern, r) for r in responses]\n",
    "    return [0.5 if match else 0.0 for match in matches]\n",
    "\n",
    "\n",
    "def count_xml(text) -> float:\n",
    "    count = 0.0\n",
    "    if text.count(\"<reasoning>\\n\") == 1:\n",
    "        count += 0.125\n",
    "    if text.count(\"\\n</reasoning>\\n\") == 1:\n",
    "        count += 0.125\n",
    "    if text.count(\"\\n<answer>\\n\") == 1:\n",
    "        count += 0.125\n",
    "        count -= len(text.split(\"\\n</answer>\\n\")[-1]) * 0.001\n",
    "    if text.count(\"\\n</answer>\") == 1:\n",
    "        count += 0.125\n",
    "        count -= (len(text.split(\"\\n</answer>\")[-1]) - 1) * 0.001\n",
    "    return count\n",
    "\n",
    "\n",
    "def xmlcount_reward_func(completions, **kwargs) -> list[float]:\n",
    "    contents = [completion[0][\"content\"] for completion in completions]\n",
    "    return [count_xml(c) for c in contents]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_prompt_length = 256\n",
    "\n",
    "from trl import GRPOConfig, GRPOTrainer\n",
    "training_args = GRPOConfig(\n",
    "    learning_rate = 5e-6,\n",
    "    adam_beta1 = 0.9,\n",
    "    adam_beta2 = 0.99,\n",
    "    weight_decay = 0.1,\n",
    "    warmup_ratio = 0.1,\n",
    "    lr_scheduler_type = \"cosine\",\n",
    "    optim = \"paged_adamw_8bit\",\n",
    "    logging_steps = 1,\n",
    "    per_device_train_batch_size = 1,\n",
    "    gradient_accumulation_steps = 1, # Increase to 4 for smoother training\n",
    "    num_generations = 6, # Decrease if out of memory\n",
    "    max_prompt_length = max_prompt_length,\n",
    "    max_completion_length = max_seq_length - max_prompt_length,\n",
    "    # num_train_epochs = 1, # Set to 1 for a full training run\n",
    "    max_steps = 250,\n",
    "    save_steps = 250,\n",
    "    max_grad_norm = 0.1,\n",
    "    report_to = \"none\", # Can use Weights & Biases\n",
    "    output_dir = \"outputs_QWen2.5-7B-vllm\",\n",
    "    # use_vllm = False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = GRPOTrainer(\n",
    "    model = model,\n",
    "    processing_class = tokenizer,\n",
    "    reward_funcs = [\n",
    "        xmlcount_reward_func,\n",
    "        soft_format_reward_func,\n",
    "        strict_format_reward_func,\n",
    "        int_reward_func,\n",
    "        correctness_reward_func,\n",
    "    ],\n",
    "    args = training_args,\n",
    "    train_dataset = dataset,\n",
    ")\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 推理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = tokenizer.apply_chat_template([\n",
    "    {\"role\" : \"user\", \"content\" : \"How many r's are in strawberry?\"},\n",
    "], tokenize = False, add_generation_prompt = True)\n",
    "\n",
    "from vllm import SamplingParams\n",
    "sampling_params = SamplingParams(\n",
    "    temperature = 0.8,\n",
    "    top_p = 0.95,\n",
    "    max_tokens = 1024,\n",
    ")\n",
    "output = model.fast_generate(\n",
    "    [text],\n",
    "    sampling_params = sampling_params,\n",
    "    lora_request = None,\n",
    ")[0].outputs[0].text\n",
    "\n",
    "output # 该输出没有推理过程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_lora(\"grpo_saved_lora\") # 保存训练好的lora权重矩阵"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = tokenizer.apply_chat_template(\n",
    "    [\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "        {\"role\": \"user\", \"content\": \"Calculate Pi\"},\n",
    "    ],\n",
    "    tokenize=False,\n",
    "    add_generation_prompt=True,\n",
    ")\n",
    "\n",
    "from vllm import SamplingParams\n",
    "\n",
    "sampling_params = SamplingParams(\n",
    "    temperature=0.8,\n",
    "    top_p=0.95,\n",
    "    max_tokens=1024,\n",
    ")\n",
    "output = (\n",
    "    model.fast_generate(\n",
    "        text,\n",
    "        sampling_params=sampling_params,\n",
    "        lora_request=model.load_lora(\"./grpo_saved_lora\"),\n",
    "    )[0]\n",
    "    .outputs[0]\n",
    "    .text\n",
    ")\n",
    "\n",
    "output # 该输出有推理过程"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
